import cv2 as cv
import npfnn
import torch
from zhnnet_npfnn import ZhnNet1
from zhnnet_pytorch import ZhnNet2
import time

print('npfnn test.')
model1 = ZhnNet1()
model2 = ZhnNet2()
# npfnn.save(model1.state_dict(), 'model1')
model1.load_state_dict(npfnn.load('model1'))
model_dict1 = npfnn.load('model1')
model_dict2 = model2.state_dict()
model_list2 = list(model_dict2.keys())
len1 = len(model_dict1)
len2 = len(model_list2)
m, n = 0, 0
while True:
    if m >= len1 or n >= len2:
        break
    layername2 = model_list2[n]
    w2 = model_dict2[layername2]
    w1 = model_dict1[m]
    if w1.shape != w2.numpy().shape:
        continue
    model_dict2[layername2] = torch.tensor(w1, dtype=torch.float32)
    m += 1
    n += 1
model2.load_state_dict(model_dict2)
image_origin = cv.imread('npfnntestimg.jpg')
assert image_origin is not None, 'Image does not exit.'
image = image_origin.transpose(2, 0, 1) / 256
image_t = torch.tensor(image, dtype=torch.float32).unsqueeze(0)
predict2 = model2(image_t)[0].detach().numpy()
# print('start.')
t1 = time.time()
predict1 = model1(image)
print(f'time consumption: {time.time() - t1}')
err = abs(predict1 - predict2)
err = sum(err.flatten())
print(f'total error: {err}')
